Word squares

Time: O(N^2xN!); Space: O(N^2); hard

Given a set of words without duplicates, find all word squares you can build from them.

A sequence of words forms a valid word square if the kth row and column read the exact same string, where 0 ≤ k < max(numRows, numColumns).

For example, the word sequence [“ball”,“area”,“lead”,“lady”] forms a word square because each word reads the same both horizontally and vertically.

b a l l
a r e a
l e a d
l a d y

Notes:

  • There are at least 1 and at most 1000 words.

  • All words will have the exact same length.

  • Word length is at least 1 and at most 5.

  • Each word contains only lowercase English alphabet a-z.

Example 1:

Input: words =

["area",
 "lead",
 "wall",
 "lady",
 "ball"
]

Output: [[“wall”,“area”,“lead”,“lady”], [“ball”,“area”,“lead”,“lady”]]

Explanation:

  • The output consists of two word squares.

  • The order of output does not matter (just the order of words in each word square matters).

Example 2:

Input: words =

["abat",
 "baba",
 "atan",
 "atal"
]

Output: [[“baba”,“abat”,“baba”,“atan”], [“baba”,“abat”,“baba”,“atal”]]

[1]:

class TrieNode(object):
    def __init__(self):
        self.indices = []
        self.children = [None] * 26

    def insert(self, words, i):
        cur = self
        for c in words[i]:
            if not cur.children[ord(c)-ord('a')]:
                cur.children[ord(c)-ord('a')] = TrieNode()
            cur = cur.children[ord(c)-ord('a')]
            cur.indices.append(i)

[2]:

class Solution1(object):
    """
    Time: O(N^2*N!)
    Space: O(N^2)
    """
    def wordSquares(self, words):
        """
        :type words: List[str]
        :rtype: List[List[str]]
        """
        result = []

        trie = TrieNode()
        for i in range(len(words)):
            trie.insert(words, i)

        curr = []
        for s in words:
            curr.append(s)
            self.wordSquaresHelper(words, trie, curr, result)
            curr.pop()

        return result

    def wordSquaresHelper(self, words, trie, curr, result):
        if len(curr) >= len(words[0]):
            return result.append(list(curr))

        node = trie
        for s in curr:
            node = node.children[ord(s[len(curr)]) - ord('a')]
            if not node:
                return

        for i in node.indices:
            curr.append(words[i])
            self.wordSquaresHelper(words, trie, curr, result)
            curr.pop()

[3]:

s = Solution1()

words = [
  "area",
  "lead",
  "wall",
  "lady",
  "ball"
]
assert s.wordSquares(words) == [["wall","area","lead","lady"], ["ball","area","lead","lady"]]

words = [
  "abat",
  "baba",
  "atan",
  "atal"
]
assert s.wordSquares(words) == [["baba","abat","baba","atan"], ["baba","abat","baba","atal"]]